module Apx (T(..), mk, fromExpr, eval, scale) where

import qualified Ply
import qualified Bs
import qualified Iv
import qualified Expr
import Interlude

data T c =
    P (Ply.T c)
  | Add (T c) (T c)
  | Mul (T c) (T c)
  | Neg (T c)
  | Pow (T c) Integer
  | Uni Expr.ElemFun (T c)

mk :: (Ord a, Floating a) => T a -> Bool -> Maybe (Ply.T a)
mk (P p)      _ = Just p
mk (Add t u)  d = liftA2 (+) (mk t d) (mk u d)

mk (Mul t u)  d = compress 3 d <$>
                  liftA2 (*) (isPly t <|> (mkt . (d==) =<< sign u mku))
                             (isPly u <|> (mku . (d==) =<< sign t mkt))
                  where mkt = cacheEnum $ mk t
                        mku = cacheEnum $ mk u

mk (Neg t) d = negate <$> mk t (not d)

mk (Pow t k)  d | k < 0          = mk (Pow (Uni Expr.Recip t) (-k)) d
mk (Pow _ 0)  _                  = Just $ Ply.const 1
mk (Pow t k)  d | k `mod` 2 == 1 = (^k) <$> mk t d
mk (Pow t k)  d                  = mk (t `Mul` Pow t (k-1)) d

mk (Uni f t)  d = (\t' -> fApx d (eval (repeat 0.5) t) (Bs.lohi $ Bs.ofPly t')
                             `Ply.compose` t') <$> mk t (fMono d)
  where (fApx, fMono) = case f of
                          Expr.Recip -> (Apx.recip_,not)
                          Expr.Exp   -> (Apx.exp_  ,id )
                          Expr.Sqrt  -> (Apx.sqrt_ ,id )
                          Expr.Atan  -> (Apx.atan_ ,id )

eval :: Floating a => [a] -> T a -> a
eval x (P p)     = Ply.eval p x
eval x (Add t u) = eval x t + eval x u
eval x (Mul t u) = eval x t * eval x u
eval x (Neg t)   = - eval x t
eval x (Pow t k) = eval x t ^ k
eval x (Uni f t) = Expr.elemFun f $ eval x t

scale :: Num c => [Iv.T c] -> T c -> T c
scale dom = sc where
  sc (P p)     = P $ Ply.scale p dom
  sc (Add t u) = Add (sc t) (sc u)
  sc (Mul t u) = Mul (sc t) (sc u)
  sc (Neg t)   = Neg (sc t)
  sc (Pow t k) = Pow (sc t) k
  sc (Uni f t) = Uni f (sc t)

fromExpr :: Floating a => Expr.T a -> T a
fromExpr (Expr.Pi)      = P $ Ply.const pi
fromExpr (Expr.Const c) = P $ Ply.const c
fromExpr (Expr.Var i)   = P $ Ply.var i (Iv.mk 0 1)
fromExpr (Expr.Add (fromExpr -> P t) (fromExpr -> P u)) = P $ t + u
fromExpr (Expr.Add (fromExpr ->   t) (fromExpr ->   u)) = Add t u
fromExpr (Expr.Mul (fromExpr -> P t) (fromExpr -> P u)) = P $ t * u
fromExpr (Expr.Mul (fromExpr ->   t) (fromExpr ->   u)) = Mul t u
fromExpr (Expr.Neg (fromExpr -> P t))   = P $ - t
fromExpr (Expr.Neg (fromExpr ->   t))   = Neg t
fromExpr (Expr.Pow (fromExpr -> P t) k) = P $ t ^ k
fromExpr (Expr.Pow (fromExpr ->   t) k) = Pow t k
fromExpr (Expr.Elem f (fromExpr -> P (Ply.isConst -> Just c))) = P $ Ply.const $ Expr.elemFun f c
fromExpr (Expr.Elem f (fromExpr -> t                        )) = Uni f t


sign :: (Floating c, Ord c) => T c -> (Bool -> Maybe (Ply.T c)) -> Maybe Bool
sign (P p) _ = Bs.sign $ Bs.ofPly p
sign t axt | eval (repeat 0.5) t >= 0 && join (Bs.sign . Bs.ofPly <$> axt True)  == Just True  = Just True
sign t axt | eval (repeat 0.5) t <= 0 && join (Bs.sign . Bs.ofPly <$> axt False) == Just False = Just False
sign _ _   = Nothing

compress :: (Ord a, Fractional a) => Int -> Bool -> Ply.T a -> Ply.T a
compress k d p = if d then p' + Ply.const rl else p' + Ply.const rh where
  p' = Ply.lo k p
  (rl,rh) = Iv.dest $ Bs.lohi $ Bs.ofPly $ p - p'

isPly :: T a -> Maybe (Ply.T a)
isPly (P p) = Just p
isPly _     = Nothing

linear :: Fractional a => (a,a) -> (a,a) -> Ply.T a
linear (x1,y1) (x2,y2) = Ply.const y1 + (x - Ply.const x1) * m where
  x = Ply.var 0 (Iv.mk 0 1)
  m = Ply.const $ (y2-y1)/(x2-x1)

taylor :: (Ord a, Floating a) => (Int -> a -> a) -> a -> Int -> Ply.T a
taylor f x0 k = sum $ (`map` [0..k]) $ \i ->
    Ply.const (f i x0 / fromIntegral (fact i)) * (x - Ply.const x0)^i
  where x = Ply.var 0 (Iv.mk 0 1)

-- ** Approximations

-- True: lower approximation
-- False: upper approximation

recip_ :: (Ord a, Floating a) => Bool -> a -> Iv.T a -> Ply.T a
recip_ True  r (Iv.dest -> (a,_)) | 0 <= a && 0 < r = taylor recipD r 1
recip_ False _ (Iv.dest -> (a,b)) | 0 <= a          = linear (a, recip a) (b, recip b)

exp_ :: (Ord a, Floating a) => Bool -> a -> Iv.T a -> Ply.T a
exp_ True  r _                  = taylor (\_ -> exp) r 1
exp_ False _ (Iv.dest -> (a,b)) = linear (a, exp a) (b, exp b)

sqrt_ :: (Ord a, Floating a) => Bool -> a -> Iv.T a -> Ply.T a
sqrt_ True  _ (Iv.dest -> (a,b)) | 0 <= a = linear (a, sqrt a) (b, sqrt b)
sqrt_ False r _                  | 0 <= r = taylor sqrtD r 1

atan_ :: (Ord a, Floating a) => Bool -> a -> Iv.T a -> Ply.T a
atan_ True  _ (Iv.dest -> (a,b)) | 0 <= a           = linear (a, atan a) (b, atan b)
atan_ False r (Iv.dest -> (a,_)) | 0 <= a && 0 <= r = taylor atanD r 1

-- ** Derivatives

recipD :: (Ord a, Fractional a) => Int -> a -> a
recipD 0 x = recip x
recipD 1 x = - recip (x*x)

sqrtD :: (Ord a, Floating a) => Int -> a -> a
sqrtD 0 x = sqrt x
sqrtD 1 x = recip (2 * sqrt x)

atanD :: Floating a => Int -> a -> a
atanD 0 x = atan x
atanD 1 x = recip (1 + x*x)
